import asyncio
import websockets
import json

REMOTE_WS_URL = 'wss://api.tenclass.net/xiaozhi/v1/'
EXTRA_HEADERS = {
    "Authorization": "Bearer test-token",
    "Device-Id": "f0:9e:9e:0e:ad:b0",
    "Protocol-Version": 3
}

connected_clients = set()
remote_ws = None

async def connect_to_remote():
    global remote_ws
    remote_ws = await websockets.connect(REMOTE_WS_URL, extra_headers=EXTRA_HEADERS)
    print("Connected to remote WebSocket server")

async def forward_to_remote(message):
    if remote_ws and remote_ws.open:
        await remote_ws.send(message)
        if isinstance(message, bytes):
            protocol = BinaryProtocol(message)
            print("Message forwarded to remote server:: ", protocol.to_dict())
        else:
            print(f"Message forwarded to remote server:: {message}")

async def broadcast_to_clients(message):
    if connected_clients:  # 如果有连接的客户端
        await asyncio.wait([client.send(message) for client in connected_clients])

async def forward_from_remote_to_clients():
    global intercept
    intercept = False
    while True:
        if remote_ws and remote_ws.open:
            try:
                message = await remote_ws.recv()
                if isinstance(message, str) and 'sentence' in message:
                    msg = json.loads(message)
                    print(msg)
                    if msg['state'] == 'sentence_start':
                        intercept = True
                        msg['text'] = '今天天气不好!所以,滚蛋!'
                        await broadcast_to_clients(json.dumps(msg, ensure_ascii=False))
                        await send_custom_msg()
                        continue
                    elif msg['state'] == 'sentence_end':
                        msg['text'] = '今天天气不好!所以,滚蛋!'
                        intercept = False
                        await broadcast_to_clients(json.dumps(msg, ensure_ascii=False))
                        continue
                if intercept:
                    continue
                
                if isinstance(message, bytes):
                    protocol = BinaryProtocol(message)
                    print("Received message from remote server: ", protocol.to_dict())
                else:
                    print(f"Received message from remote server: {message}")
                await broadcast_to_clients(message)
            except websockets.exceptions.ConnectionClosed:
                print("Remote server connection closed")
                break
        else:
            await asyncio.sleep(1)  # 每秒检查一次连接状态

async def handler(websocket, path):
    connected_clients.add(websocket)
    try:
        async for message in websocket:
            if isinstance(message, bytes):
                protocol = BinaryProtocol(message)
                print("Received message from client: ", protocol.to_dict())
            else:
                print(f"Received message from client: {message}")
            
            await forward_to_remote(message)
    except websockets.exceptions.ConnectionClosed as e:
        print("A client just disconnected.")
    finally:
        connected_clients.remove(websocket)

async def send_custom_msg():
    from audio_processing import read_pcm_file, split_into_frames, encoder
    from binary_protocol import NewBinaryProtocol

    # 读取 PCM 文件
    pcm_file_path = './wav/gun.pcm'
    pcm_data = read_pcm_file(pcm_file_path)

    # 将 PCM 数据分割成帧
    frames = split_into_frames(pcm_data, FRAME_SIZE)

    for frame in frames:
        # 编码 PCM 帧为 Opus 数据
        opus_frame = encoder.encode(frame.tobytes(), FRAME_SIZE)

        # 序列化 BinaryProtocol
        audio_message = NewBinaryProtocol(type=0, reserved=0, payload=opus_frame)
        audio_serialized = audio_message.serialize()

        # 发送 Opus 数据
        await broadcast_to_clients(audio_serialized)
        print(f"Sent frame of size {len(opus_frame)} bytes")